import cv2
import numpy as np

# 加载图像
image = cv2.imread("input.jpg")
height, width = image.shape[:2]

# 将图像转换为HSV色彩空间
hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

# 设置白色背景的HSV范围
lower_color = np.array([0, 0, 200])  # 白色下限
upper_color = np.array([180, 30, 255])  # 白色上限

# 创建掩码，将白色背景区域提取出来
color_mask = cv2.inRange(hsv_image, lower_color, upper_color)

# 创建一个边缘掩码，仅保留图像的边缘区域
margin = 50  # 边缘宽度，可以根据需要调整
center_mask = np.zeros((height, width), dtype=np.uint8)
cv2.rectangle(center_mask, (margin, margin), (width - margin, height - margin), 255, -1)
edge_mask = cv2.bitwise_not(center_mask)

# 将颜色掩码与边缘掩码相结合
combined_mask = cv2.bitwise_and(color_mask, edge_mask)

# 反转掩码，将前景提取出来
foreground = cv2.bitwise_and(image, image, mask=cv2.bitwise_not(combined_mask))

# 创建新背景（选择红、蓝或白色背景）
background_color = {
    'blue': [255, 0, 0],
    'red': [0, 0, 255],
    'white': [255, 255, 255]
}

# 选择所需的背景颜色，例如蓝色背景
bg_color = background_color['blue']

# 创建一个纯色背景，并与前景合成
background = np.full(image.shape, bg_color, dtype=np.uint8)
background = cv2.bitwise_and(background, background, mask=combined_mask)
result = cv2.add(foreground, background)

# 显示结果
cv2.imshow("Result", result)
cv2.waitKey(0)
cv2.destroyAllWindows()

# 保存结果图像
cv2.imwrite("output.jpg", result)
